
# _*_ coding:utf-8 _*_
import os
import xml.etree.ElementTree as ET
from tqdm import tqdm

dirpath = r'datasets/ciwa_zonghe/Annotations'  # 原来存放xml文件的目录


total_classes = ['_biaozhunbengque', '_biaozhunliewen', '_lunkuobengque', '_qianmo', 
    '_tudian', '_zazhi', '_duocengmian', '_xixiaoliewen', '_daojiao', '_jiejing', 
    '_mohubengque', '_bailiewen']

qx_class = []
qx_num = {}

i = 0
names = []
FileNames = os.listdir(dirpath)

save_list = []

for fp in tqdm(FileNames):
    if fp.endswith('.xml'):
        # i = i + 1
        root = ET.parse(os.path.join(dirpath, fp)).getroot()

        xmin, ymin, xmax, ymax = 0, 0, 0, 0
        sz = root.find('size')
        width = float(sz[0].text)
        height = float(sz[1].text)

        # if width > max_w:
        #     max_w = width
        # if height > max_h:
        #     max_h = height

        filename = root.find('filename').text
        num_obj = 0
        for child in root.findall('object'):  # 找到图片中的所有框

            name = child.findtext('name')
            if name not in total_classes:
                continue
            box = child.findall('robndbox')
            if name not in qx_class:
                qx_class.append(name)
                qx_num[name] = 1
    
            else:
                qx_num[name] += 1
            num_obj += 1
        
        if num_obj > 0:
            save_list.append(filename)
        

with open('trainval.txt', 'w') as f:
    for i in save_list:
        f.write(i)
        f.write('\n')
f.close()

print('缺陷类别有：', qx_class)
print('各缺陷个数为：', qx_num)
# print(qx_class)
# print("max_h: ", max_h)
# print("max_w: ", max_w)